package com.aries.dubbo.like.rpc.service.util.discover.zk;

import com.aries.dubbo.like.rpc.common.util.ClassUtil;
import lombok.extern.slf4j.Slf4j;
import org.I0Itec.zkclient.ZkClient;
import org.apache.zookeeper.data.Stat;

import java.net.Inet4Address;
import java.net.UnknownHostException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Created with IntelliJ IDEA.
 * Author: aries
 * Date: 2018/8/24
 * Description: 服务注册和更新。
 */
@Slf4j
public class ZookeeperRegistry {
    private static final String BASE_PATH = "/aries-rpc";
    private static Map<String, List<String>> localRegisterData;

    private static void doRegister(Map<String, List<String>> map) {
        localRegisterData = map;
        log.info("do register.");
        ZkClient zkClient = ZookeeperInstance.getZkClient();
        if (!zkClient.exists(BASE_PATH)) {
            zkClient.createPersistent(BASE_PATH);
        }
        Map<String, List<String>> data = zkClient.readData(BASE_PATH, new Stat());
        if (data != null) {
            map.forEach((key, value) -> {
                if (data.containsKey(key)) {
                    List<String> dataValue = data.get(key);
                    dataValue.addAll(value);
                    List<String> newData = dataValue.stream().distinct().collect(Collectors.toList());
                    data.put(key, newData);
                } else {
                    data.put(key, value);
                }
            });
            zkClient.writeData(BASE_PATH, data);
        } else {
            zkClient.writeData(BASE_PATH, map);
        }
    }

    /**
     * 注册服务，在rpc server启动时调用。
     *
     * @param port
     */
    public static void register(int port) {
        Set<Class<?>> classSet = ClassUtil.getRpcProviderClassSet();
        Map<String, List<String>> map = new HashMap<>();
        try {
            String localIp = Inet4Address.getLocalHost().getHostAddress();
            classSet.forEach(clz -> {
                LinkedList<String> ips = new LinkedList<>();
                ips.add(localIp + ":" + port);
                map.put(clz.getInterfaces()[0].getName(), ips);
                log.info("interface " + clz.getInterfaces()[0].getName() + " will be exported on localhost:" + port);
            });
        } catch (UnknownHostException e) {
            throw new RuntimeException("error to get local ip address", e);
        }
        doRegister(map);
    }


    public static String getBasePath() {
        return BASE_PATH;
    }

    public static Map<String, List<String>> getLocalRegisterData() {
        return localRegisterData;
    }

    /**
     * 如果客户端发现某一条链路不通，就把zk注册信息里的这条链路移除,更新zk及本地信息。
     *
     * @param address 不通链路的地址
     */
    public static void updateZkData(String address) {
        Map<String, List<String>> zkData =
                ZookeeperInstance.getZkClient().readData(BASE_PATH, new Stat());
        zkData.forEach((k, v) -> {
            v.remove(address);
        });
        log.info("address:" + address + " will be remove from zk.");
        localRegisterData = zkData;
        ZookeeperInstance.getZkClient().writeData(BASE_PATH, zkData);
    }
}
